from transformers import Trainer, TrainingArguments
from src.rag import RAG
from src.data_preprocessing import preprocess_pdf
from safetensors.torch import load_file
import torch

def train_model():
    # 初始化RAG模型
    rag = RAG()
    
    # 加载预处理数据（需替换实际PDF路径）
    papers = ['data/papers/paper1.pdf', 'data/papers/paper2.pdf']
    for pdf_path in papers:
        text, embeddings = preprocess_pdf(pdf_path)
        # 将embeddings存入向量数据库
        
    # 配置训练参数
    # 使用分块保存策略（适合大模型）
    training_args = TrainingArguments(
        save_safetensors=True,
        shard_size="2GB"  # 每个分片最大2GB
    )
    
    # 启动训练（需补充实际训练数据加载逻辑）
    # trainer = Trainer(model=rag.model, args=training_args)
    # trainer.train()

def load_safe_checkpoint(path):
    state_dict = load_file(path)
    model.load_state_dict(state_dict)
    
    # 在训练过程中添加以下代码
    # 安全加载训练检查点
    # load_safe_checkpoint("checkpoint.safetensors")
    
    # 传统方式保存（训练过程中）
    # torch.save(model.state_dict(), "checkpoint.pth") 
    
    # 最终产出保存为安全格式
    # model.save_pretrained("./final_model", safe_serialization=True)

if __name__ == "__main__":
    train_model()
